# -------------------------------------------------------------------------------
# Filename:     export_testlog2excel.py
# UpdateDate:   2024/06/08
# Description:  一键导出测试日志中 Average(3iters)行的关键信息到Excel文件中。
# Example:      python export_testlog2excel.py
# Depends:
# Notes:        本脚本适用于Docker容器环境下运行；
# -------------------------------------------------------------------------------
import os
import re
import sys
import glob
import pandas as pd
from os.path import basename
from openpyxl import load_workbook
from openpyxl.styles import Font, Border, Side, PatternFill, Alignment
from openpyxl.utils import get_column_letter

#如果报错，可安装必要的库（如果尚未安装）：
# pip install pandas openpyxl

# 获取当前脚本文件的目录
current_dir = os.path.dirname(os.path.abspath(__file__))
# 获取上一级目录
parent_dir = os.path.dirname(current_dir)
# 构造上一级目录下的log子目录的完整路径
log_directory = os.path.join(parent_dir, 'log')
# 输出的Excel文件
output_excel = log_directory + '/report.xlsx'
print(f"log_directory: {log_directory}")
print(f"output_excel: {output_excel}")
#sys.exit(1)
# 正则表达式匹配行（根据需要调整）
#test_regex = re.compile(r'Average\(+12+iters\)\s+([0-9]+\.?[0-9]*)\s+([0-9]+\.?[0-9]*)\s+([0-9]+\.?[0-9]*)\s+([0-9]+\.?[0-9]*)\s+([0-9]+\.?[0-9]*)\s+([0-9]+\.?[0-9]*)\s+([0-9]+\.?[0-9]*)')
#test_regex = re.compile(r'Average\(+[0-9]+iters\)\s+([0-9]+\.?[0-9]*)\s+([0-9]+\.?[0-9]*)\s+([0-9]+\.?[0-9]*)\s+([0-9]+\.?[0-9]*)\s+([0-9]+\.?[0-9]*)\s+([0-9]+\.?[0-9]*)\s+([0-9]+\.?[0-9]*)')
test_regex = re.compile(r'Average\(+[0-9]+iters\)\s+([0-9]+\.?[0-9]*)\s+([0-9]+\.?[0-9]*)\s+([0-9]+\.?[0-9]*)\s+([0-9]+\.?[0-9]*)\s+([0-9]+\.?[0-9]*)\s+([0-9]+\.?[0-9]*)\s+([0-9]+\.?[0-9]*)\s+([0-9]+\.?[0-9]*)')

log_files = glob.glob(os.path.join(log_directory, '*.log'))
if not log_files:
    print(f"Error: The log file does not exist in {log_directory}")
    sys.exit(1)

# 初始化列表来存储数据
test_data = []
# 有效数据条数的计数器, {count_data+2}来对应表格的行数。
count_data = 0
# 1. 提取关键信息到Excel表格中： 查找目录下的所有.log文件中关键信息，保存到Excel表格中
for log_file in glob.glob(os.path.join(log_directory, '*.log')):
    #print(f"log_file： {log_file}")
    # 1.1. 解析固定参数
    pattern = r'_blks(\d+)_sli(\d+)_slo(\d+)_tp(\d+)_bs(\d+)_'
    match = re.search(pattern, log_file)
    if match:
        blks_number = match.group(1)
        sli_number = match.group(2)
        slo_number = match.group(3)
        tp_number = match.group(4)
        #print("提取的数字是:", blks_number, sli_number, slo_number, tp_number)
    # 1.2. 解析动态参数
    with open(log_file, 'r') as file:
        for line in file:
            test_match = test_regex.match(line)
            if test_match:
                #print(f"log_file: {log_file}")
                file_name = basename(log_file)  # 提取文件名
                print(f"file_name: {file_name}")
                # 使用os.path.splitext()去除后缀，返回元组(文件名, 后缀)
                filename_without_extension, extension = os.path.splitext(file_name)
                #print("test_match:", int(test_match.group(1)), int(test_match.group(2)), int(test_match.group(3)), int(test_match.group(4)))
                # 提取数据并添加到列表中
                test_data.append({
                    'tp_number': int(tp_number),
                    'block_size': int(blks_number),
                    'input_seqlen': int(sli_number),
                    'output_seqlen': int(slo_number),
                    'batch_size': int(test_match.group(1)),
                    'context_latency(ms)': float(test_match.group(2)),
                    'per_token_latency(ms)': float(test_match.group(3)),
                    'context_latency_device(ms)': float(test_match.group(4)),
                    'per_token_latency_device(ms)': float(test_match.group(5)),
                    'e2e_latency(ms)': float(test_match.group(6)),
                    'e2e_throughput(tokens/s)': float(test_match.group(7)),
                    'decoder_throughput(tokens/s)': float(test_match.group(8)),
                    'log_test': file_name,
                    'log_test_cnmon': filename_without_extension + "_cnmon.log"
                })
                # 增加有效数据条数的计数器
                count_data += 1
    #print("test_data:", test_data)
    # 1.3. 保存到Excel中。将列表转换为DataFrame
    df_test = pd.DataFrame(test_data)
    with pd.ExcelWriter(output_excel) as writer:
        df_test.to_excel(writer, sheet_name='TestReport')
# 2. 排序Excel文件：对已经保存的Excel文件进行排除，方便查看
df = pd.read_excel(output_excel)
df_sorted = df.sort_values(by=['tp_number', 'block_size', 'input_seqlen', 'output_seqlen', 'batch_size'], ascending=[True, True, True, True, True])
df_sorted.to_excel(output_excel, index=False)

# 3. 整理新序号列
df = pd.read_excel(output_excel)
df = df.iloc[:, 1:]
df.reset_index(drop=True, inplace=True)
df.insert(0, 'Number', range(1, len(df) + 1))
df.to_excel(output_excel, index=False)

#print("有效数据总条数:", count_data)
print(f"Excel汇总报告已保存到： {output_excel}")

# 4. 设置表格格式
# 加载已存在的Excel文件
workbook = load_workbook(output_excel)
worksheet = workbook.active  # 或者通过名称获取特定的sheet: worksheet = workbook['Sheet1']
#worksheet = workbook['TestReport']

#worksheet.column_dimensions['A'].width = 20
# 4.1. 设置列宽
for col in 'ABCDEF':
    worksheet.column_dimensions[col].width = 8
for col in 'GHIJKLMNO':
    worksheet.column_dimensions[col].width = 16

# 4.2. 设置第一行的行高（设置为35磅）
worksheet.row_dimensions[1].height = 35

# 4.3. 设置文本自动换行样式
#wrap_text_alignment = Alignment(wrap_text=True)
# 设置文本自动换行、左对齐并垂直居中的样式
wrap_text_alignment = Alignment(wrap_text=True, horizontal='left', vertical='center')

# 遍历第一行的所有单元格并设置自动换行
for col in range(1, 16):  # 假设第一行有数据到第12列（L列）
    cell = worksheet.cell(row=1, column=col)
    cell.alignment = wrap_text_alignment  # 设置自动换行

# 4.4. 设置边框样式
thin_border = Border(left=Side(style='thin'),
                     right=Side(style='thin'),
                     top=Side(style='thin'),
                     bottom=Side(style='thin'))

# 4.5. 设置浅蓝色和浅绿色填充样式
light_blue_fill = PatternFill(start_color='ACC5E1', end_color='ACC5E1', fill_type='solid')
light_green_fill = PatternFill(start_color='BDFCC9', end_color='BDFCC9', fill_type='solid')
# 遍历前14列与第2到 {count_data+2} 行的范围
for row in range(1, count_data+2):  # 注意行号从1开始
    for col in range(1, 16):  # 前15列，即A到O
        cell = worksheet.cell(row=row, column=col)
        # 设置边框
        cell.border = thin_border
        # 设置背景色
        if row == 1:  # 第一行（在循环中是第2行）设置为浅蓝色
            cell.fill = light_blue_fill
        # 设置日志超链接，确保单元格有值
        if ( col == 14 or col == 15) and row > 1:
            #print(f"====cell.value: {cell.value}")
            if cell.value :
                #print(f"====cell.value: {cell.value}")
                #cell.value = basename(cell.value) # 获取文件名并更新单元格值
                cell.hyperlink = "./" + cell.value  # 设置超级链接为当前单元格的值（即文件的完整路径）
                font = Font(color="0000FF", underline="single")  # 创建一个Font对象并设置其属性
                cell.font = font  # 直接将Font对象赋值给单元格的font属性
                # 重新设置单元格值
                if col == 14 :
                    cell.value = "test.log"
                if col == 15 :
                    cell.value = "test_cnmon.log"

# 4.6. 保存更改后的Excel文件
workbook.save(output_excel)

print(f"Excel表格式已设置完毕： {output_excel}")

